#######################
###Plot the estimates##
#######################
library(ggplot2)
library(dplyr)
med_plots <- function(model.name, sim_results_list, output_dir = "~/OneDrive - Indiana University/FromGoogle/syd/plots/") {
  # Create output directory if it doesn't exist
  if (!dir.exists(output_dir)) {
    dir.create(output_dir, recursive = TRUE)
  }
  
  setwd(output_dir)
  print(paste("Output directory set to:", output_dir))
  
  # Define labels
  area_labels <- c("Agricultural", "Forest", "Other")
  outcome_means <- c(
    mean(med.dat.ag$clear_zoon_confirm, na.rm = TRUE),
    mean(med.dat.for$clear_zoon_confirm, na.rm = TRUE),
    mean(med.dat.other$clear_zoon_confirm, na.rm = TRUE)
  )
  
  # Summary function
  compute_summary <- function(values) {
    mean_val <- mean(values)
    lower_ci <- quantile(values, 0.025)
    upper_ci <- quantile(values, 0.975)
    return(c(Mean = mean_val, Lower = lower_ci, Upper = upper_ci))
  }
  
  # Build summary table
  summary_list <- mapply(function(sim_results, area_label) {
    ACME <- unlist(lapply(sim_results, function(x) x$d.avg.sims))
    ADE <- unlist(lapply(sim_results, function(x) x$z.avg.sims))
    Total <- unlist(lapply(sim_results, function(x) x$tau.sims))
    
    summary_matrix <- t(sapply(list(ACME, ADE, Total), compute_summary))
    colnames(summary_matrix) <- c("Mean", "Lower", "Upper")
    
    df <- data.frame(
      Area = area_label,
      Effect = c("ACME", "ADE", "Total Effect"),
      summary_matrix
    )
    return(df)
  }, sim_results_list, area_labels, SIMPLIFY = FALSE)
  
  summary_all <- do.call(rbind, summary_list)
  summary_all$Effect <- factor(summary_all$Effect, levels = c("ACME", "ADE", "Total Effect"))
  summary_all$Area <- factor(summary_all$Area, levels = c("Agricultural", "Forest", "Other"))
  summary_all$row_order <- with(summary_all, interaction(Effect, Area, drop = TRUE))
  summary_all$row_order <- factor(summary_all$row_order, levels = c(
    "Total Effect.Agricultural", "Total Effect.Forest", "Total Effect.Other",
    "ADE.Agricultural", "ADE.Forest", "ADE.Other",
    "ACME.Agricultural", "ACME.Forest", "ACME.Other"
  ))
  
  effect_colors <- c("ACME" = "darkgreen", "ADE" = "orange", "Total Effect" = "brown")
  divider_pos <- c(3.5, 6.5)
  
  # ==== PLOT 1 ====
  print("Generating Plot 1: Effect Sizes")
  jpeg(filename = paste0("fig_combined_effects_", model.name, ".jpeg"), width = 8, height = 6, units = 'in', res = 500)
  print(
    ggplot(summary_all, aes(y = Mean, x = row_order, color = Effect)) +
      geom_point() +
      geom_errorbar(aes(ymin = Lower, ymax = Upper), width = 0.2) +
      geom_hline(yintercept = 0, linetype = "dashed", color = "black", size = 1) +
      geom_vline(xintercept = divider_pos, linetype = "dotted", color = "grey40", size = 0.6) +
      scale_color_manual(values = effect_colors) +
      scale_x_discrete(labels = rep(c("Agricultural", "Forest", "Other"), 3)) +
      scale_y_continuous(
        limits = c(min(summary_all$Lower) * 1.1, max(summary_all$Upper) * 1.1),
        expand = c(0, 0)
      ) +
      coord_flip() +
      labs(y = "Effect size", x = "", color = "Effect type") +
      theme_minimal() +
      theme(
        panel.border = element_rect(colour = "black", fill = NA, size = 1),
        plot.title = element_text(hjust = 0.5)
      )
  )
  dev.off()
  print("Plot 1 saved.")
  
  # ==== PLOT 2 ====
  print("Generating Plot 2: Percent Change")
  summary_percent <- summary_all
  summary_percent$OutcomeMean <- rep(outcome_means, each = 3)
  summary_percent$Mean <- (summary_percent$Mean / summary_percent$OutcomeMean) * 100
  summary_percent$Lower <- (summary_percent$Lower / summary_percent$OutcomeMean) * 100
  summary_percent$Upper <- (summary_percent$Upper / summary_percent$OutcomeMean) * 100
  
  jpeg(filename = paste0("fig_combined_percent_", model.name, ".jpeg"), width = 8, height = 6, units = 'in', res = 500)
  print(
    ggplot(summary_percent, aes(y = Mean, x = row_order, color = Effect)) +
      geom_point() +
      geom_errorbar(aes(ymin = Lower, ymax = Upper), width = 0.2) +
      geom_hline(yintercept = 0, linetype = "dashed", color = "black", size = 1) +
      geom_vline(xintercept = divider_pos, linetype = "dotted", color = "grey40", size = 0.6) +
      scale_color_manual(values = effect_colors) +
      scale_x_discrete(labels = rep(c("Agricultural", "Forest", "Other"), 3)) +
      scale_y_continuous(
        limits = c(min(summary_percent$Lower) * 1.1, max(summary_percent$Upper) * 1.1),
        expand = c(0, 0),
        labels = scales::percent_format(scale = 1)
      ) +
      coord_flip() +
      labs(y = "Effect size (%)", x = "", color = "Effect type") +
      theme_minimal() +
      theme(
        panel.border = element_rect(colour = "black", fill = NA, size = 1),
        plot.title = element_text(hjust = 0.5)
      )
  )
  dev.off()
  print("Plot 2 saved.")
}

